"""
文件名：example2.py

在客船在指定码头间的定向旅客运输时，不同承运公司的不同的船班，对旅客的吸引能力时不同的。
假设船班的对旅客的吸引力主要体现在如下几方面：
- （1）船班时间，time_hour，分别为 0,1,2,..., 23.
- （2）船的大小，size，分别为 1,2,3.
- （3）所属船公司的差异，company，分别为1,2,3 分别代表承运公司A、承运公司B、其他承运公司。
- （4）当天的所有船公司的 船总数量 n_ship，范围是 1,2,3,...,20

另外已知的数据为
- （5）当天的历史旅客总数。sum_passenger
- （6）船班的旅客数。n_passenger

预测的目标值y值为
- （1）船班分走的旅客比例；ratio

"""
import copy
import datetime
import math
import random
from typing import Tuple, NoReturn, Dict, List

import pandas

from profiler.profiler import MultiParamProfiler


# step0，数据仿真器

def func(x_dict: Dict[str, float], param_dict: Dict[str, float]) -> float:
    """
    拟合函数
    :param x_dict: time_hour, size, company, n_ship
    :param param_dict: good_time_factor, company_{1,2,3}_factor, size_factor, normalize_factor, n_ship_factor
    :return: 吸引旅客比例
    """
    time_hour: int = round(x_dict.get("time_hour"))
    size: int = round(x_dict.get("size"))
    company: int = round(x_dict.get("company"))
    n_ship: int = round(x_dict.get("n_ship"))

    good_time_factor: float = param_dict.get("good_time_factor") if time_hour in {7, 8, 9, 10, 18, 19, 20, 21} else 0.
    company_factor = param_dict.get(f"company_{company}_factor")
    size_factor = param_dict.get("size_factor")
    normalize_factor = param_dict.get("normalize_factor")
    n_ship_factor = math.log(n_ship) / math.log(2. + param_dict.get("n_ship_factor"))

    ratio = normalize_factor * (good_time_factor + size * size_factor + company_factor + n_ship_factor)
    return ratio


_param_dict: Dict[str, float] = {
    "good_time_factor": 5,
    "company_1_factor": 0.2,
    "company_2_factor": 0.6,
    "company_3_factor": 0.8,
    "size_factor": 1.5,
    "n_ship_factor": 1.,
    "normalize_factor": 1,
}


def _simulate() -> Tuple[pandas.DataFrame, pandas.DataFrame]:
    n_point = 100  # 点数
    time_hours = tuple(range(0, 24))  # 船班时间
    sizes = (1, 2, 3)  # 船大小
    companies = (1, 2, 3)  # 服务等级
    n_ships = tuple(range(1, 20 + 1))  # 总船数
    sum_passenger_range = (5000, 10000)  # 总旅客数
    best_x = {"time_hour": 8, "size": sizes[-1], "company": companies[-1], "n_ship": n_ships[-1], }
    _param_dict["normalize_factor"] = 1. / func(x_dict=best_x, param_dict=_param_dict, )
    x_dict_list: List[Dict[str, float]] = []
    y_dict_list: List[Dict[str, float]] = []
    for i in range(n_point):
        x_dict: Dict[str, float] = {
            "time_hour": random.choice(time_hours),
            "size": random.choice(sizes),
            "company": random.choice(companies),
            "n_ship": random.choice(n_ships),
        }
        sum_passenger = random.randint(*sum_passenger_range)
        ratio_ideal = func(x_dict=x_dict, param_dict=_param_dict)
        n_passenger_ideal = max(1, round(sum_passenger * ratio_ideal))
        n_passenger = random.gauss(mu=float(n_passenger_ideal), sigma=math.sqrt(float(n_passenger_ideal)))
        ratio = float(n_passenger) / float(sum_passenger)
        error_raito = math.sqrt(float(n_passenger)) / float(sum_passenger)
        y_dict: Dict[str, float] = {
            "ratio": ratio,
            "error_ratio": error_raito,
        }
        x_dict_list.append(copy.copy(x_dict))
        y_dict_list.append(copy.copy(y_dict))
    df_x: pandas.DataFrame = pandas.DataFrame(x_dict_list)  # 列名为 time_hour, size, level, n_ship
    df_y: pandas.DataFrame = pandas.DataFrame(y_dict_list)  # 列名为 ratio，error_ratio
    return df_x, df_y


def _fit_and_test():
    # 拟合数据、测试数据拆分
    df_x, df_y = _simulate()
    n_row, _ = df_x.shape
    n_test = 20
    n_fit = n_row - n_test
    df_x_fit: pandas.DataFrame = df_x.iloc[:n_fit, :]
    df_y_fit: pandas.DataFrame = df_y.iloc[:n_fit, :]
    df_x_test: pandas.DataFrame = df_x.iloc[n_fit:, :]
    df_y_test: pandas.DataFrame = df_y.iloc[n_fit:, :]

    print(f"\n-- df_x_fit, shape: {df_x_fit.shape}, head: \n{df_x_fit.head().to_string()}")
    print(f"\n-- df_y_fit, shape: {df_y_fit.shape}, head: \n{df_y_fit.head().to_string()}")
    print(f"\n-- df_x_test, shape: {df_x_test.shape}, head: \n{df_x_test.head().to_string()}")
    print(f"\n-- df_y_test, shape: {df_y_test.shape}, head: \n{df_y_test.head().to_string()}")

    # 参数初始值与范围
    param_info_dict_list = [
        {"name": "good_time_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
        {"name": "company_1_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
        {"name": "company_2_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
        {"name": "company_3_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
        {"name": "size_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
        {"name": "n_ship_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
        {"name": "normalize_factor", "init": 1.e-5, "lower": 0., "upper": 10.},
    ]
    df_param_input = pandas.DataFrame(param_info_dict_list)  # 参数设定值dataframe
    print(f"\n-- df_param_input, shape: {df_param_input.shape}, content: \n{df_param_input.to_string()}")

    # 拟合
    time_fit_start = datetime.datetime.now()
    model: MultiParamProfiler = MultiParamProfiler()
    df_param_fitted = model.fit(df_x=df_x_fit, df_y=df_y_fit, df_param=df_param_input,
                                func=func, y_col="ratio", error_y_col="error_ratio")
    time_fit_end = datetime.datetime.now()
    cost_time_fit = time_fit_end - time_fit_start
    print(f"-- fit cost time seconds: {cost_time_fit.total_seconds()}")

    # 预测
    x_dict_test_list: List[Dict[str, float]] = df_x_test.to_dict("records")
    y_dict_test_list: List[Dict[str, float]] = df_y_test.to_dict("records")
    y_predict_info_dict_list: List[Dict[str, float]] = []
    for x_dict_test, y_dict_test in zip(x_dict_test_list, y_dict_test_list):
        y_test = y_dict_test.get("ratio")
        y_predict, _ = model.predict(x_dict=x_dict_test, func_error=None)
        y_predict_info_dict_list.append(copy.copy({
            "ratio_predict": y_predict,
            "ratio_simulation": y_test,
            "diff_absolute": y_predict - y_test,
            "diff_relative": None if y_test == 0 else (y_predict - y_test) / y_test,
        }))
    df_y_predict_info = pandas.DataFrame(y_predict_info_dict_list)
    print(f"\n-- df_y_predict_info, shape: {df_y_predict_info.shape}, content: \n{df_y_predict_info.to_string()}")

    # 参数对比
    param_comparison_dict_list: List[Dict[str, float | str]] = df_param_fitted.to_dict("records")
    for o in param_comparison_dict_list:
        o["set"] = _param_dict.get(o.get("name"))
    df_param_comparison = pandas.DataFrame(param_comparison_dict_list)
    print(f"\n-- df_param comparison, shape: {df_param_comparison.shape}, content: \n{df_param_comparison.to_string()}")


if __name__ == '__main__':
    _fit_and_test()
